#include <stdio.h>
#include <iostream>
#include <string>

#include <pcl/io/ply_io.h>
#include <pcl/io/pcd_io.h>
#include <pcl/point_types.h>
#include <pcl/io/pcd_io.h>
#include <pcl/io/ply_io.h>
#include <pcl/visualization/cloud_viewer.h>
#include <pcl/filters/uniform_sampling.h>
#include <pcl/features/normal_3d_omp.h>
#include <pcl/features/shot_omp.h>
#include <pcl/features/vfh.h>
#include <pcl/common/common.h>
#include <pcl/filters/voxel_grid.h>

#include <fstream>
#include <sys/types.h>
#include <sys/stat.h>

#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/ml/ml.hpp>

typedef pcl::PointXYZ PointType;
class classifier{
    private:
        cv::Ptr<cv::ml::SVM> SVMclf;
    public:
        classifier(void){
            SVMclf = cv::ml::SVM::create();
        }
        classifier(std::string filename){
            SVMclf = cv::ml::SVM::create();
            SVMclf->load(filename.c_str());
        }
        void trainSVM(cv::Mat _trainingData, cv::Mat _trainingLabels){
            SVMclf->setType(cv::ml::SVM::Types::C_SVC);
            SVMclf->setKernel(cv::ml::SVM::KernelTypes::LINEAR);
            SVMclf->setTermCriteria(cv::TermCriteria(cv::TermCriteria::MAX_ITER, 100, 1e-6));

            cv::Ptr<cv::ml::TrainData> tData = cv::ml::TrainData::create(_trainingData, cv::ml::ROW_SAMPLE,_trainingLabels);
            SVMclf->train(tData);
            std::cout << "SVM trained " << std::endl;
            SVMclf->save("SVMclf->xml");
        }
        void trainSVM(cv::Mat _trainingData, cv::Mat _trainingLabels,std::string fileName){
            SVMclf->setType(cv::ml::SVM::Types::C_SVC);
            SVMclf->setKernel(cv::ml::SVM::KernelTypes::LINEAR);
            SVMclf->setTermCriteria(cv::TermCriteria(cv::TermCriteria::MAX_ITER, 100, 1e-6));

            cv::Ptr<cv::ml::TrainData> tData = cv::ml::TrainData::create(_trainingData, cv::ml::ROW_SAMPLE,_trainingLabels);
            SVMclf->train(tData);
            std::cout << "SVM trained " << std::endl;
            fileName+=".xml";
            SVMclf->save(fileName.c_str());
        }

        float getConfidence(float distance){
            float conf = 1.0 / ( 1.0 + exp(-1*distance));
            return conf;
        }
        void validateSVM(cv::Mat _testData, cv::Mat _testLabels){
            float hits = 0;
            float miss = 0;
            for(size_t idx=0;idx<_testData.rows;idx++){
                std::cout << " [test] "<< "idx:"<< idx << " " << SVMclf->predict(_testData.row(idx)) << " "
                        << _testLabels.at<float>(idx,0) << " [ confidence ] "
                        << getConfidence(SVMclf->predict(_testData.row(idx))) <<  std::endl;
                if (getConfidence(SVMclf->predict(_testData.row(idx)))==_testLabels.at<float>(idx,0))
                    hits++;
                else
                    miss++;
            }
            printf(" [accuracy] %f ",(hits/(hits+miss)));
        }
        std::vector<float> predict(cv::Mat query){
            std::vector<float> result;
            result.push_back( SVMclf->predict(query));
            result.push_back( getConfidence(SVMclf->predict(query)));
            return result;
        }
        cv::Mat vector2Mat(pcl::PointCloud<pcl::VFHSignature308>::Ptr inputDescriptor){
            cv::Mat testArray = cv::Mat::zeros(1, 306, CV_32FC1);
            for(size_t j=0; j<306;j++){
                testArray.at<float>(0,j)= (float)inputDescriptor->points[0].histogram[j];
            }
            return testArray;
        }
};

class processPointCloud{
    private:
        float model_ss_;
    public:
        pcl::PointCloud<PointType>::Ptr cloud;
        pcl::PointCloud<PointType>::Ptr cloud_keypoints;
        pcl::PointCloud<pcl::Normal>::Ptr cloud_normals;
        pcl::PointCloud<pcl::VFHSignature308>::Ptr VFH_descriptor;

        pcl::UniformSampling<PointType> uniform_sampling;
        pcl::NormalEstimationOMP<PointType, pcl::Normal> norm_est;

        processPointCloud(void){
            cloud = pcl::PointCloud<PointType>::Ptr(new pcl::PointCloud<PointType>);
            cloud_keypoints = pcl::PointCloud<PointType>::Ptr(new pcl::PointCloud<PointType>);
            cloud_normals = pcl::PointCloud<pcl::Normal>::Ptr(new pcl::PointCloud<pcl::Normal>);
            VFH_descriptor = pcl::PointCloud<pcl::VFHSignature308>::Ptr(new pcl::PointCloud<pcl::VFHSignature308>);
            model_ss_= 0.001f;
        }

        void cloudinput(pcl::PointCloud<PointType>::Ptr someCloud){
            cloud = someCloud;
        }

        void getNormals(void){
            norm_est.setKSearch (10);
            norm_est.setInputCloud (cloud);
            norm_est.compute (*cloud_normals);
        }

        void getKeypoints(void){
            pcl::UniformSampling<PointType> uniform_sampling;
            uniform_sampling.setInputCloud (cloud);
            uniform_sampling.setRadiusSearch (model_ss_);
            uniform_sampling.filter(*cloud_keypoints);
        }

        void getVFHE(void){
            pcl::VFHEstimation<PointType, pcl::Normal, pcl::VFHSignature308> vfh;
            pcl::search::KdTree<PointType>::Ptr kdtree(new pcl::search::KdTree<PointType>);
            vfh.setInputCloud(cloud);
            vfh.setInputNormals(cloud_normals);
            vfh.setSearchMethod(kdtree);
            vfh.setNormalizeBins(true);
            vfh.setNormalizeDistance(false);
            vfh.compute(*VFH_descriptor);
        }

        pcl::PointCloud<pcl::VFHSignature308>::Ptr getDescriptor(void){
            getNormals();
            //getKeypoints();
            getVFHE();
            return VFH_descriptor;
        }
};

class getPointCloud{
    private :
        struct stat info;
    public :
        std::string PCD_BASE_PATH;
        std::string PCD_CLASS_PATH;
        std::string PCD_FILE_PATH;
        std::string PCD_FILE_EXT;
        std::string FULL_FILE_PATH;
        std::string CLASSIFIER_NAME;
        std::vector<std::string> DataFiles;
        pcl::PointCloud<PointType>::Ptr cloud;
        pcl::PointCloud<pcl::VFHSignature308>::Ptr VFH_descriptor;

        ofstream outputFile;

        processPointCloud pc;
        classifier clf;

        getPointCloud(std::string _PCD_BASE_PATH,std::vector<std::string> _DataFiles,std::string _CLASSIFIER_NAME){
            cloud = pcl::PointCloud<PointType>::Ptr(new pcl::PointCloud<PointType>);
            VFH_descriptor = pcl::PointCloud<pcl::VFHSignature308>::Ptr(new pcl::PointCloud<pcl::VFHSignature308>);
            PCD_BASE_PATH = _PCD_BASE_PATH;
            CLASSIFIER_NAME = _CLASSIFIER_NAME;
            for(size_t i=0;i<_DataFiles.size();i++){
                DataFiles.push_back(PCD_BASE_PATH +_DataFiles[i]);
            }
        }

        pcl::PointCloud<PointType>::Ptr loadCloud(std::string filePath){
            pcl::PointCloud<PointType>::Ptr someCloud (new pcl::PointCloud<PointType>);
            if(filePath.find(".ply")==filePath.npos){
                if (pcl::io::loadPCDFile<PointType> (filePath, *someCloud) == -1){ //* load the file
                    PCL_ERROR ("Couldn't read file test_pcd.pcd \n");
                }
            }
            else{
                if (pcl::io::loadPLYFile<PointType> (filePath, *someCloud) == -1){ //* load the file
                    PCL_ERROR ("Couldn't read file test_ply.ply \n");
                }
            }

            return someCloud;
        }

        void getTrainingData(std::vector <cv::Mat>& _trainingData,std::vector <cv::Mat>& _trainingLabels){
            float minX,minY,minZ=1000.0;
            float maxX,maxY,maxZ=0.0;

            pcl::VoxelGrid<PointType> vg;
            cv::Mat label = cv::Mat::zeros(1,1,CV_32SC1);

            // Loop over each Directory
            for(size_t fileIdx=0; fileIdx<DataFiles.size();fileIdx++){
                std::vector<std::string> files = std::vector<std::string>();
                PointType minPt, maxPt;
                int skip=0;

                getdir(DataFiles[fileIdx]+"/",files);
                printf("No of Files: %d\r\n",(int)files.size());

                cv::Mat trainingData;
                cv::Mat trainingLabels;
                // Loop over each file
                for (unsigned int i = 0;i < files.size();i++) {
                    if(files[i] == "."  || files[i] == ".."){
                        printf(" . or .. files ignored\r\n");
                    }
                    else{
                        std::cout << " [ processing ] " << files[i];

                        cloud = loadCloud(DataFiles[fileIdx]+"/"+files[i]);

                        //-- voxelization of the cloud --------//
                        vg.setInputCloud (cloud);
                        vg.setLeafSize (0.005f, 0.005f, 0.005f);
                        vg.filter (*cloud);
                        //-- Getting the descriptors ----------//
                        pc.cloudinput(cloud);
                        VFH_descriptor = pc.getDescriptor();
                        cv::Mat _descriptor = cv::Mat::zeros(1,306,CV_32FC1);

                        for(size_t i=0;i<306;i++){
                            _descriptor.at<float>(0,i)=(float)VFH_descriptor->points[0].histogram[i];
                        }
                        //-------------------------------------//
                        trainingData.push_back(_descriptor);
                        if(DataFiles[fileIdx] == (PCD_BASE_PATH+CLASSIFIER_NAME)){
                            label.at<float>(0,0)=1.0;
                            trainingLabels.push_back(label);
                            std::cout << " [label] 1.0";
                        }
                        else{
                            label.at<float>(0,0)= -1.0;
                            trainingLabels.push_back(label);
                            std::cout << " [label] -1.0";
                        }
                        std::cout << std::endl;

                    }
                }
                _trainingData.push_back(trainingData);
                _trainingLabels.push_back(trainingLabels);
            }
        }
        int getdir (std::string dir, std::vector<std::string> &files){
            DIR *dp;
            struct dirent *dirp;
            if((dp  = opendir(dir.c_str())) == NULL) {
                std::cout << "Error(" << errno << ") opening " << dir << endl;
                return errno;
            }

            while ((dirp = readdir(dp)) != NULL) {
                if(std::string(dirp->d_name)!="."&&std::string(dirp->d_name)!="..")
                    files.push_back(std::string(dirp->d_name));
            }
            closedir(dp);
            return 0;
        }

};

void write2file(std::string DataFilePath,std::string Fname ,cv::Mat Data){
        cv::FileStorage file(DataFilePath, cv::FileStorage::WRITE);
        file << Fname << Data;
        file.release();
}

void readFromFile(std::string DataFilePath,std::string Fname ,cv::Mat &Data){
    cv::FileStorage file(DataFilePath, cv::FileStorage::READ);
    file[Fname] >> Data;
    file.release();
}

int main(int argc, char** argv){
    std::string pkgPath = ".";
    PCL_INFO("Training SVM\r\n");
    std::string PCD_BASE_PATH = "./assets/";
    std::string POSITIVE_CLASS = "zhuantou";
    bool parsePCD = true,trainSVM = true;
    std::vector<std::string> PCD_CLASS_PATHS;
    PCD_CLASS_PATHS.push_back("bowl");
    PCD_CLASS_PATHS.push_back("cap");
    PCD_CLASS_PATHS.push_back("cereal_box");
    PCD_CLASS_PATHS.push_back("coffee_mug");
    PCD_CLASS_PATHS.push_back("coffee_table");
    PCD_CLASS_PATHS.push_back("office_chair");
    PCD_CLASS_PATHS.push_back("zhuantou");
    PCD_CLASS_PATHS.push_back("soda_can");
    PCD_CLASS_PATHS.push_back("sofa");
    PCD_CLASS_PATHS.push_back("table");
    PCL_INFO("Parameters %s %s\r\n",PCD_BASE_PATH.c_str(),parsePCD ? "true":"false");

    for(size_t i=0;i<PCD_CLASS_PATHS.size();i++){
        PCL_INFO("class: %s\r\n",PCD_CLASS_PATHS[i].c_str());
    }
    PCL_INFO("done\r\n");
    std::string trainDataFile = pkgPath+"/bin/Dataset/";
    std::string trainLabelsFile = pkgPath+"/bin/Dataset/trainingLabels";
    std::string classifierPath = pkgPath+"/bin/"+POSITIVE_CLASS;
    PCL_INFO("trainData: %s\r\n",trainDataFile.c_str());
    PCL_INFO("trainLabels: %s\r\n",trainLabelsFile.c_str());

    classifier clf;

    if(parsePCD){
        std::vector< cv::Mat > _trainingData;
        std::vector< cv::Mat > _trainingLabels;

        getPointCloud p(PCD_BASE_PATH,PCD_CLASS_PATHS,POSITIVE_CLASS);
        p.getTrainingData(_trainingData,_trainingLabels);

        int ret = mkdir("./bin/Dataset", 0755);
        for (size_t f=0;f<_trainingData.size();f++){
            PCL_INFO("file: %s\r\n",PCD_CLASS_PATHS[f].c_str());
            PCL_INFO("Training Data dimensions: rows: %d cols: %d\r\n",(int)_trainingData[f].rows,(int)_trainingData[f].cols);
            PCL_INFO("Training Labels dimensions: rows: %d cols: %d\r\n",(int)_trainingLabels[f].rows,(int)_trainingLabels[f].cols);
            std::stringstream idx;
            idx << f;
            write2file(trainDataFile+PCD_CLASS_PATHS[f]+".xml",PCD_CLASS_PATHS[f],_trainingData[f]);
        }

        PCL_INFO("Training Data saved.\r\n");
    }

    if(trainSVM){
        cv::Mat trainingData;
        cv::Mat trainingLabels;


        cv::Mat tempLabels = cv::Mat::zeros(1,1,CV_32SC1);
        float labelval=-1.0;

        for(size_t f=0;f< PCD_CLASS_PATHS.size();f++){
            cv::Mat temptrainingData;
            readFromFile(trainDataFile+PCD_CLASS_PATHS[f]+".xml",PCD_CLASS_PATHS[f],temptrainingData);
            trainingData.push_back(temptrainingData);

            if (PCD_CLASS_PATHS[f] == POSITIVE_CLASS){
                labelval=1.0;
            }
            else{
                labelval= -1.0;
            }

            for(size_t d=0;d<temptrainingData.rows;d++){
                tempLabels.at<float>(0,0)=labelval;
                trainingLabels.push_back(tempLabels);
            }
            PCL_INFO("[ CLASS] %s [ LABEL] %f\r\n",PCD_CLASS_PATHS[f].c_str(),labelval);
        }


        PCL_INFO("Files Loaded.\r\n");
        PCL_INFO("training SVM.\r\n");
        clf.trainSVM(trainingData,trainingLabels,classifierPath);
        clf.validateSVM(trainingData,trainingLabels);
        PCL_INFO("SVM trained.\r\n");
    }

    return 0;
}
